package com.agent.rag.client;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Component
public class SplitterClient {
    private Long minChunksSize=512L; //文本分割块大小
    private float maxCosineSimilarity=0.7f; //相似度极限
    private final EmbeddingModel embeddingModel; //嵌入模型

    @Autowired
    public SplitterClient(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    /**
     * 设置分割块大小
     * @param minChunksSize
     * @return
     */
    public SplitterClient setChunksSize(Long minChunksSize) {
        this.minChunksSize = minChunksSize;
        return this;
    }

    /**
     *  设置相似度界限
     * @param maxCosineSimilarity
     * @return
     */
    public SplitterClient setSimilarity(Long maxCosineSimilarity){
        this.maxCosineSimilarity = maxCosineSimilarity;
        return this;
    }

    /**
     * 文本切割
     * @param text
     * @return
     */
    public List<String> spiltText(String text){
        List<String> sentences=splitSentences(text);
        List<float[]> embeddings=new ArrayList<>();
        for(String sentence:sentences){
            embeddings.add(embeddingModel.embed(sentence));
        }
        //获取文本边界
        List<Integer> bounds=getBounds(embeddings);
        return merge(sentences,bounds);
    }

    /**
     * 切割单句文本合并
     * @param sentences
     * @param bounds
     * @return
     */
    private List<String> merge(List<String> sentences,List<Integer> bounds){
        StringBuilder sentenceBuilder=new StringBuilder();
        List<String> chunks=new ArrayList<>();
        Long currentSize=0L;
        for(int i=0;i<sentences.size();i++){
            //文本块边界判断
            if(bounds.contains(i)){
                //添加文本块
                chunks.add(sentenceBuilder.toString().trim());
                //初始化成新文本块的起点
                sentenceBuilder.setLength(0);
                sentenceBuilder.append(sentences.get(i));
                //设置新文本块长度
                currentSize+=sentences.get(i).length();
                //如果边界值也是集合的最后一个元素，直接添加文本块
                if(i==sentences.size()-1){
                    chunks.add(sentenceBuilder.toString().trim());
                }
                //跳转到下一次循环
                continue;
            }
            sentenceBuilder.append(sentences.get(i));
            currentSize+=sentences.get(i).length();
            //对文本块超界判断
            if(i==sentences.size()-1||currentSize>minChunksSize){
                chunks.add(sentenceBuilder.toString().trim());
                sentenceBuilder.setLength(0);
                currentSize=0L;
            }
        }
        return chunks;
    }

    /**
     * 获取句子的边界
     * @param sentences
     * @return
     */
    private List<Integer> getBounds(List<float[]> sentences){
        float cosineSimilarity;
        List<Integer> bounds=new ArrayList<>();
        for(int i=1;i<sentences.size();i++){
            cosineSimilarity=cosineSimilarity(sentences.get(i-1),sentences.get(i));
            if(cosineSimilarity<maxCosineSimilarity){
                bounds.add(i);
            }
        }
        return bounds;
    }

    /**
     * 文本句子切分
     * @param text
     * @return
     */
    private List<String> splitSentences(String text){
        return List.of(text.split("(?<=[。！？])"));
    }

    /**
     * 余弦相似度计算
     * @param a
     * @param b
     * @return
     */
    private float cosineSimilarity(float[] a,float[] b){
        float vecAll=0.0f;
        float vecOne=0.0f;
        float vecTwo=0.0f;
        for(int i=0;i<a.length;i++){
            vecAll+=a[i]*b[i];
            vecOne+=a[i]*a[i];
            vecTwo+=b[i]*b[i];
        }
        float sqrtOne= (float) Math.sqrt(vecOne);
        float sqrtTwo= (float) Math.sqrt(vecTwo);
        return vecAll/(sqrtOne*sqrtTwo);
    }
}
